package com.hadluo.qq.server.config;

import com.alibaba.fastjson.JSON;
import com.hadluo.qq.server.entity.Message;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.*;

import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Component
public class ChatWebSocketHandler implements WebSocketHandler {
    // 用户会话映射表，key为用户ID，value为WebSocket会话
    private static final Map<Long, WebSocketSession> USER_SESSIONS = new ConcurrentHashMap<>();

    @Override
    public void afterConnectionEstablished(WebSocketSession session) {
        // 从session中获取用户ID
        Long userId = getUserId(session);
        if (userId != null) {
            USER_SESSIONS.put(userId, session);
        }
    }

    @Override
    public void handleMessage(WebSocketSession session, WebSocketMessage<?> message) throws IOException {
        // 解析消息
        String payload = message.getPayload().toString();
        Message chatMessage = JSON.parseObject(payload, Message.class);

        // 获取接收者的会话
        WebSocketSession receiverSession = USER_SESSIONS.get(chatMessage.getReceiverId());
        if (receiverSession != null && receiverSession.isOpen()) {
            // 发送消息给接收者
            receiverSession.sendMessage(new TextMessage(payload));
        }
    }

    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) {
        // 处理传输错误
        Long userId = getUserId(session);
        if (userId != null) {
            USER_SESSIONS.remove(userId);
        }
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus closeStatus) {
        // 连接关闭后清理会话
        Long userId = getUserId(session);
        if (userId != null) {
            USER_SESSIONS.remove(userId);
        }
    }

    @Override
    public boolean supportsPartialMessages() {
        return false;
    }

    /**
     * 从WebSocket会话中获取用户ID
     */
    private Long getUserId(WebSocketSession session) {
        try {
            String userId = session.getAttributes().get("userId").toString();
            return Long.parseLong(userId);
        } catch (Exception e) {
            return null;
        }
    }

    /**
     * 向指定用户发送消息
     */
    public void sendMessageToUser(Long userId, Message message) {
        WebSocketSession session = USER_SESSIONS.get(userId);
        if (session != null && session.isOpen()) {
            try {
                session.sendMessage(new TextMessage(JSON.toJSONString(message)));
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }

    /**
     * 获取用户的WebSocket会话
     */
    public WebSocketSession getSession(Long userId) {
        return USER_SESSIONS.get(userId);
    }

    /**
     * 判断用户是否在线
     */
    public boolean isUserOnline(Long userId) {
        WebSocketSession session = USER_SESSIONS.get(userId);
        return session != null && session.isOpen();
    }
}